from grid_env_ideal_obs_repeat_task import *
from grid_agent import *
from checkpoint_utils import *
from maze_factory import *
from replay_config import *
import argparse
import json
import sys
import matplotlib.pyplot as plt
from sklearn.manifold import TSNE
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Circle
from matplotlib.lines import Line2D
from sklearn.manifold import TSNE
import random
from sklearn.decomposition import PCA
from matplotlib.animation import FuncAnimation
from sklearn.cluster import KMeans
import threading
import mplcursors
from mpl_toolkits.mplot3d.art3d import Poly3DCollection



# 載入 "./logs/max_le_collection.npy" 和 "./logs/max_ler_collection.npy" 這兩個矩陣

max_le_collection = np.load("./logs/max_le_collection.npy")
max_ler_collection = np.load("./logs/max_ler_collection.npy")

print("shape of max_le_collection: ", max_le_collection.shape)
print("shape of max_ler_collection: ", max_ler_collection.shape)

# 將 max_le_collection 和 max_ler_collection 中的所有記錄繪製出來
fig = plt.figure()
for i in range(max_le_collection.shape[0]):
    # 用每一条线的 i 值添加标签
    plt.plot(max_ler_collection[i], label=str(i)+'_ler')
    plt.plot(max_le_collection[i], label=str(i)+'_le')
plt.legend()
plt.show()

# 繪製 max_le_collection 和 max_ler_collection 各自的包絡

max_ler_collection_upper_bound = np.max(max_ler_collection, axis=0)
max_le_collection_upper_bound = np.max(max_le_collection, axis=0)
max_ler_collection_lower_bound = np.min(max_ler_collection, axis=0)
max_le_collection_lower_bound = np.min(max_le_collection, axis=0)

epsilons = [(i+1)*0.01 for i in range(max_le_collection.shape[1])]

fig = plt.figure()
plt.fill_between(epsilons, max_ler_collection_lower_bound, max_ler_collection_upper_bound, color='red', alpha=0.1)
plt.fill_between(epsilons,  max_le_collection_lower_bound, max_le_collection_upper_bound, color='blue', alpha=0.1)
# plot mean value
plt.plot(epsilons, np.mean(max_ler_collection, axis=0), color='red', label='PLE-random')
plt.plot(epsilons, np.mean(max_le_collection, axis=0), color='blue', label='PLE-trained')

# 横轴标注 epsilon, 纵轴标注 primary LE
plt.xlabel('epsilon')
plt.ylabel('primary LE')
plt.legend()
plt.show()